#!/usr/bin/python
# lrwiz
"""
Lib-Ray Wizard script to create a Lib-Ray volume.
"""
# Initially just a collection of useful functions

# Requires the "avconv" command-line utility.

import os, subprocess, getopt, string
from userinput import *
#from chapter import Chapter, ChapterSet

import avprobe

TEMPDIR = "."
            
            
standard_resolutions = {
    'sd':       (720,   405),
    'hd720':    (1280,  720),
    'hd1080':   (1920,  1080),
    '4k':       (4096,  2304)
    }
standard_resolution_names = ['sd', 'hd720', 'hd1080', '4k']
            
            
class VideoSource(object):
    """
    Information for a video stream to be included in the Lib-Ray Video file.
    """
    
    # Shell commands required to implement conversions
    cmd_src_to_y4m_w_rate = "avconv -i {SRCVID!s} -s {RESOLUTION!s} -r {FRAMERATE!s} -f yuv4mpegpipe -pix_fmt yuv420p {Y4MVID!s}"
    
    cmd_src_to_y4m_wo_rate = "avconv -i {SRCVID!s} -s {RESOLUTION!s} -f yuv4mpegpipe -pix_fmt yuv420p {Y4MVID!s}"
    
    cmd_y4m_to_vp8_w_rate = "vpxenc {Y4MVID!s} -o {VP8VID!s} --i420 -w {WIDTH:d} -h {HEIGHT:d} -p 2 -t 3 --best --target-bitrate={BITRATE:d} --end-usage=vbr --auto-alt-ref=1 --fps={FRAMERATE:f} -v --minsection-pct=5 --maxsection-pct=800 --lag-in-frames=16 --kf-min-dist=0 --kf-max-dist=480 --token-parts=3 --static-thresh=0 --drop-frame=0 --min-q=0 --max-q=6"
    
    cmd_y4m_to_vp8_wo_rate = "vpxenc {Y4MVID!s} -o {VP8VID!s} --i420 -w {WIDTH:d} -h {HEIGHT:d} -p 2 -t 3 --best --target-bitrate={BITRATE:d} --end-usage=vbr --auto-alt-ref=1 --minsection-pct=5 --maxsection-pct=800 --lag-in-frames=16 --kf-min-dist=0 --kf-max-dist=480 --token-parts=3 --static-thresh=0 --drop-frame=0 --min-q=0 --max-q=6"
    
    cmd_vp8_to_y4m = "avconv -i {VP8VID!s} -f yuv4mpegpipe {Y4RVID!s}"
    
    cmd_psnr = "qpsnr -a avg_psnr -o fpa={FPA:d} -r {Y4MVID!s} {Y4RVID!s} > {PSNR!s}"
    
    cmd_ssim = "qpsnr -a avg_ssim -o fpa={FPA:d} -r {Y4MVID!s} {Y4RVID!s} > {SSIM!s}"
    
    def __init__(self, num):
        self.num = num
        self.imagestream = False
        self.path = ""
        self.name = "Video%2.2d" % num
        self._override_framerate = False
        self.framerate = 30
        self.use_mux_audio = False
        self.parameters = {}
        
    def __str__(self):
        return "%s=%s (Images? %s, Framerate=%s, Use Audio? %s)" % (
                self.name, os.path.basename(self.path), 
                str(self.imagestream), str(self.framerate), str(self.use_mux_audio))
                
    def _get_filepaths(self):
        filepaths = {
            'SRCVID':   self.path,
            'Y4MVID':   os.path.join(TEMPDIR, self.name + '.y4m'),
            'Y4RVID':   os.path.join(TEMPDIR, self.name + '.y4r'),
            'VP8VID':   os.path.join(TEMPDIR, self.name + '.vp8'),
            'PSNR':     os.path.join(TEMPDIR, self.name + '.psnr'),
            'SSIM':     os.path.join(TEMPDIR, self.name + '.ssim')
            }
        return filepaths
            
        
    def script(self, resolution, bitrate=20000, keep=False, metrics=True):
        script = "\n# Ingest script for Video Source %s (%s)\n" % (self.path, self.name)
        script += self.src_to_y4m(resolution)
        script += self.y4m_to_vp8(resolution, bitrate)
        if metrics:
            script += "# Computing quality metrics\n"
            script += self.vp8_to_y4m()
            script += self.compute_metrics()
        if not keep:
            script += self.clean_up()
        return script
        
    def clean_up(self):
        parameters = self._get_filepaths()
        return "rm {Y4MVID!s} {Y4RVID!s}".format(**parameters) + '\n'
        
    def src_to_y4m(self, resolution):
        parameters = {
                     'RESOLUTION': resolution,
                     'FRAMERATE': self.framerate
                     }      
        parameters.updat(self._get_filepaths())
        if self.framerate:
            return self.cmd_src_to_y4m_w_rate.format(**parameters) + '\n'
        else:
            return self.cmd_src_to_y4m_wo_rate.format(**parameters) + '\n'
            
    def y4m_to_vp8(self, resolution, bitrate=20000):
        if resolution in standard_resolutions:
            width, height = standard_resolutions[resolution]
        else:
            raise ValueError("Unsupported resolution %s." % resolution)
        
        parameters = {
            'RESOLUTION': resolution,
            'WIDTH': width,
            'HEIGHT': height,
            'BITRATE': bitrate,
            'FRAMERATE': self.framerate
            }
        parameters.update(self._get_filepaths())
        if self.framerate and self._override_framerate:
            return self.cmd_y4m_to_vp8_w_rate.format(**parameters) + '\n'
        else:
            return self.cmd_y4m_to_vp8_wo_rate.format(**parameters) + '\n'
            
    def vp8_to_y4m(self):
        parameters = self._get_filepaths()
        return self.cmd_vp8_to_y4m.format(**parameters) + '\n'
        
    def compute_metrics(self):
        if self.framerate:
            fpa = self.framerate*60
        else:
            fpa = 1800 # Assume 30fps if we don't know framerate -- really should 
                       # query the video file for this info 
        parameters = {
            'FPA': int(fpa),
            }
        parameters.update(self._get_filepaths())
        return (self.cmd_psnr.format(**parameters)  + '\n' + 
                self.cmd_ssim.format(**parameters) + '\n')
            
    def probe(self):
        """
        Use avprobe to collect basic information about the video source.
        """
        self.info = avprobe.Info(self.path)
        if self.info.framerate and not self.framerate:
            self.framerate = self.info.framerate
        # Probably want to collect other info too, such as WxH and length of video?
      
    def interactive(self):
        """
        Interactively query user for video stream information.
        
        >>> 
        >>> v = VideoSource(0)
        >>> v.interactive()
        Is this a stream of images in a directory? [y/N] y
        Location of image directory? [] dummypngs
        Frame rate for images? [30] 
        Mnemonic name for video stream after encoding? [dummypngs] 
        >>> 
        >>> str(v)
        'dummypngs=dummypngs (Images? True, Framerate=30, Use Audio? False)'
        >>> 
        >>> v = VideoSource(1)
        >>> v.interactive()
        Is this a stream of images in a directory? [y/N] 
        Video source file path? [] dummy.mp4
        Do you want to override the frame rate? [y/N] 
        Also use audio track(s) from this source file? [y/N] y
        Mnemonic name for video stream after encoding? [dummy] 
        >>> 
        >>> str(v)
        'dummy=dummy.mp4 (Images? False, Framerate=None, Use Audio? True)'
        >>>
        """
        self.imagestream = userinput(
            "Is this a stream of images in a directory?",
            self.imagestream,
            yn_inp, yn_out)
            
        if self.imagestream:
            self.path = userinput(
                "Location of image directory?", 
                default=self.path,
                inp=file_inp)
            self.framerate = userinput(
                "Frame rate for images?", 
                default=self.framerate, inp=float)
        else:
            self.path = userinput(
                "Video source file path?", 
                default=self.path,
                inp=file_inp)
            self._override_framerate = userinput(
                "Do you want to override the frame rate?",
                default=self._override_framerate, 
                inp=yn_inp, out=yn_out)
            if self._override_framerate:
                self.framerate = userinput(
                        "Framerate?", 
                        default = self.framerate,
                        inp=float)
            else:
                self.framerate = None
            self.use_mux_audio = userinput(
                "Also use audio track(s) from this source file?",
                default = self.use_mux_audio,
                inp=yn_inp, out=yn_out)
        self.name = os.path.splitext(os.path.basename(self.path))[0]                
        self.name = userinput(
                "Mnemonic name for video stream after encoding?",
                default = self.name, inp=ident_inp)

                
                
                
SOUNDTRACK, COMMENTARY = range(2)
class AudioSource(object):
    """
    Information for an audio stream to be included in the Lib-Ray Video file.
    """
    cmd_conv_to_flac = "avconv -i {SRCAUD!s} -f flac -acodec flac {FLACAUD!s}"
    
    cmd_conv_to_vorbis = "avconv -i {SRCAUD!s} -f ogg -acodec vorbis {VORBAUD!s}"
     
    cmd_copy_flac    = "cp {SRCAUD!s} {FLACAUD!s}"
    
    cmd_copy_vorbis  = "cp {SRCAUD!s} {VORBAUD!s}"
      
      
    def __init__(self, num):
        self.num = num
        self.path = ""
        self.name = "Audio%3.3d" % num
        self.language = "und"
        self.country = "ZZ"
        self.use  = SOUNDTRACK
          
    def __str__(self):
        return "%s=%s (Language: %s, Country: %s, Use: %s)" % (
              self.name, os.path.basename(self.path), 
              self.language, self.country, ["Soundtrack", "Commentary"][self.use])
              
    def _get_filepaths(self):
        """
        Get the filepaths for audio source.
        """
        filepaths = {
            'SRCAUD': self.path,
            'FLACAUD': os.path.join(TEMPDIR, self.name + '.flac'),
            'VORBAUD': os.path.join(TEMPDIR, self.name + '.ogg')
            }
        return filepaths
                
    def probe(self):
        """
        Probe source file with avprobe to check stream info.
        """
        self.info = avprobe.Info(self.path)
        if self.info.codec_name:
            self.codec_name = self.info.codec_name
            
    def script(self):
        """
        Generate ingest script for audio.
        """
        script = "\n# Ingest script for Audio Source %s (%s)\n" % (self.path, self.name)
        if self.use==SOUNDTRACK:
            script += self.src_to_flac() + '\n'
        else:
            script += self.src_to_vorbis() + '\n'
        return script
      
    def src_to_flac(self):
        """
        Convert source to FLAC using avconv (unless already provided).
        """
        parameters = self._get_filepaths()
        if self.codec_name == 'flac':
            if self.path = parameters['FLACAUD']:
                return "# No audio conversion needed for %s" % self.path
            else:
                return self.cmd_copy_flac.format(**parameters)
        else:
            return self.cmd_conv_to_flac.format(**parameters)
            
    def src_to_vorbis(self):
        """
        Convert source to Ogg Vorbis using avconv (unless already provided).
        """
        parameters = self._get_filepaths()
        if self.codec_name = 'vorbis':
            if self.path = parameters['VORBAUD']:
                return "# No audio conversion needed for %s" % self.path
            else:
                return self.cmd_copy_vorbis.format(**parameters)
        else:
            return self.cmd_conv_to_vorbis.format(**parameters)
            
          
    def interactive(self, video_source=None):
        """
        Interactively query user for information.
        
        >>> 
        >>> a = AudioSource(0)
        >>> a.interactive()
        Audio source file path? [] dummy.flac
        Language code for audio track? [und (Undetermined)] eng
        Country code for audio track? [ZZ (Universal)] 
        Use/type of audio track?
          Options:
        1 Soundtrack
        2 Commentary
        ? [Soundtrack]1
        Mnemonic name for audio track in stream? [dummy] 
        >>> 
        >>> str(a)
        'dummy=dummy.flac (Language: eng, Country: ZZ, Use: Soundtrack)'
         >>>          
        """
        if not video_source:
            self.path = userinput(
                "Audio source file path?",
                default = self.path,
                inp=file_inp)
            self.name = os.path.splitext(os.path.basename(self.path))[0]
        else:
            self.path = video_source.path
        self.language = userinput(
            "Language code for audio track?",
            default = self.language,
            inp=language_inp, out=language_out)
        self.country = userinput(
            "Country code for audio track?",
            default = self.country,
            inp=country_inp, out=country_out)
        self.use = userselect(
            "Use/type of audio track?",
            ("Soundtrack", "Commentary"),
            default = self.use)
        self.name = userinput(
            "Mnemonic name for audio track in stream?",
            default = self.name, inp=ident_inp)
                   

SUBTITLE, NOTES, CC = range(3)
class TextSource(object):
      """
      Information for a text stream to be included in the Lib-Ray Video file.
      """
      def __init__(self, num):
          self.num = num
          self.path = ""
          self.name = "Text%3.3d" % num
          self.language = "und"
          self.country = "ZZ"
          self.use  = SUBTITLE
          self.forced = False
          
      def __str__(self):
          return "%s=%s (Language: %s, Country: %s, Use: %s)" % (
                self.name, os.path.basename(self.path), 
                self.language, self.country, 
                ["Subtitles", "Screen Notes", "Closed Captioning"][self.use])
                
      def script(self):
          """
          Convert and/or generate script for ingest.
          """
          return sub_to_assa(self)
          
      def sub_to_assa(self):
          """
          Convert subtitles in formats recognized by aeidon/gaupol into ASSA format.
          """
          assa_path = os.path.join(TEMPDIR, self.name = ".assa")
          if self.path == assa_path:
              return "# Subtitles already provided in ASSA format.\n"
          else:
              sub2assa.sub2assa(self.path, assa_path)
              return "# ASSA subtitles converted %s \n" % self.assa_path 
          
      def interactive(self):
          """
          Interactively query user for information.
          
          >>> 
          >>> t = TextSource(0)
          >>> t.interactive()
          Text stream source file path? [] dummy.srt
          Language code for text track? [und (Undetermined)] fr
          Country code for text track? [ZZ (Universal)] FR
          Use/type of text track?
            Options:
          1 Subtitle
          2 Screen Notes
          3 Closed Captioning
          ? [Subtitle]
          Mnemonic name for audio track in stream? [dummy] 
          >>> 
          >>> str(t)
          'dummy=dummy.srt (Language: fr, Country: FR, Use: Subtitles)'
          >>>
          """
          self.path = userinput(
              "Text stream source file path?",
              default = self.path,
              inp=file_inp)
          self.name = os.path.splitext(os.path.basename(self.path))[0]
          self.language = userinput(
              "Language code for text track?",
              default = self.language,
              inp=language_inp, out=language_out)
          self.country = userinput(
              "Country code for text track?",
              default = self.country,
              inp=country_inp, out=country_out)
          self.use = userselect(
              "Use/type of text track?",
              ("Subtitle", "Screen Notes", "Closed Captioning"),
              default = self.use)
          self.name = userinput(
              "Mnemonic name for audio track in stream?",
              default = self.name, inp=ident_inp)
           

  
class   OutputTarget(object):
      """
      Information about the output multiplexed target file we want to create.
      """
      def __init__(self):
          self.path = ""
          self.resolution = 'hd1080'
          self._force_framerate = False
          self.framerate  = 30
          
      def __str__(self):
          return "Output=%s (Resolution=%s, Framerate=%s)" % (
              os.path.basename(self.path), str(self.resolution), str(self.framerate))

      def interactive(self, path='feature.lrv'):
          """
          Interactively query the user about the output result wanted.
          
          
          >>> 
          >>> o = OutputTarget()
          >>> o.interactive()
          Path to which to write Lib-Ray Video file? [feature.lrv] 
          Resolution desired?
            Options:
          1 Standard Definition (sd):       720 x  405
          2 High Definition (hd720):       1280 x  720
          3 Full HD (hd1080):              1920 x 1080
          4 Digital Cinema (4k):           4096 x 2304
          ? [Full HD (hd1080):              1920 x 1080]
          Do you wish to force the framerate to a specific value? [y/N] 
          >>> 
          >>> str(o)
          'Output=feature.lrv (Resolution=hd1080, Framerate=None)'
          >>>
          """
          self.path = path
          self.path = userinput(
              "Path to which to write Lib-Ray Video file?",
              default = self.path,
              inp=new_file_inp)
          self.resolution = standard_resolution_names[userselect(
              "Resolution desired?",
              ( "Standard Definition (sd):       720 x  405",
                "High Definition (hd720):       1280 x  720",
                "Full HD (hd1080):              1920 x 1080",
                "Digital Cinema (4k):           4096 x 2304"),
              default = standard_resolution_names.index(self.resolution))]
          self._force_framerate = userinput(
              "Do you wish to force the framerate to a specific value?",
              default = self._force_framerate,
              inp=yn_inp, out=yn_out)
          if self._force_framerate:
              self.framerate = userinput(
                  "Desired framerate?",
                  default = self.framerate,
                  inp=int)
          else:
              self.framerate = None
              
              
              
class   ModeSettings(object):
      """
      Information about miscellaneous operating mode settings.
      """
      def __init__(self):
          self.keep_intermediates = False
          self.compute_quality_metrics = True
          
      def __str__(self):
          return "Settings: Intermediates? %s, Compute Quality Metrics? %s" % (
              self.keep_intermediates, self.compute_quality_metrics)
      
      def interactive(self):
          """
          Query user for mode settings.
          
          >>> 
          >>> m = ModeSettings()
          >>> m.interactive()
          Do you want to keep intermediate files generated during conversion? [y/N] y
          Compute visual quality metrics for quality control report? [Y/n] 
          >>> 
          >>> str(m)
          'Settings: Intermediates? True, Compute Quality Metrics? True'
          >>>
          """
          self.keep_intermediates = userinput(
              "Do you want to keep intermediate files generated during conversion?",
              default = self.keep_intermediates,
              inp=yn_inp, out=yn_out)
          self.compute_quality_metrics = userinput(
              "Compute visual quality metrics for quality control report?",
              default = self.compute_quality_metrics,
              inp=yn_inp, out=yn_out)
              

class   LibRayVideoTitle(object):
      """
      Information needed for conversion (collects all settings).
      """
      def __init__(self):
          self.mode_settings = ModeSettings()
          self.target = OutputTarget()
          self.video_sources = []
          self.audio_sources = []
          self.text_sources = []
          self.chapters = []
          self._menu_step = 0
          
      def interactive(self):
        while True:
          self._menu_step = 0
          self._menu_step = userselect(
            "Configure Lib-Ray Video Conversion:",
            ("Add/edit video sources",
             "Add/edit audio sources",
             "Add/edit text (subtitle) Sources",
             "Add/edit chapter",
             "Configure target Lib-Ray Video settings",
             "Configure mode settings",
             "See what the wizard will do",
             "Start the conversion process",
             "Save configuration to file",
             "Load configuration from file",
             "Exit the program (without running conversion)"),
             default = self._menu_step)
          # Dispatch by number
          self._menu_step = [
                self._i_add_video, 
                self._i_add_audio, 
                self._i_add_text,
                self._i_add_chapter,
                self._i_config_target,
                self._i_config_mode,
                self.show_settings,
                self.start_conversion,
                self.save_settings,
                self.load_settings,
                self._exit
                ][self._menu_step]()
                
      def _i_add_video(self):
          i = self._i_add(self.video_sources, "Video Sources", VideoSource)
          mux_sources = {}
          for video_source in self.video_sources:
              if video_source.use_mux_audio:
                  mux_sources[video_source.path]=video_source
          for audio_source in self.audio_sources:
              if audio_source.path in mux_sources:
                  del mux_sources[audio_source.path]
          for mux in mux_sources:
                  self.audio_sources.append(AudioSource(len(self.audio_sources)))
                  self.audio_sources[-1].interactive(mux_sources[mux])
          return 2
      
      def _i_add_audio(self):
          self._i_add(self.audio_sources, "Audio Sources", AudioSource)
          return 3
          
      def _i_add_text(self):
          self._i_add(self.text_sources, "Text (Subtitle) Sources", TextSource)
          return 4
          
      def _i_add_chapter(self):
          self._i_add(self.chapters, "Chapters", Chapter)
          return 5
          
      def _i_add(self, sources, caption, constructor):
          """
          Interactively add or edit sources (or chapters).
          """
          add_more = True
          while add_more:
              options = [str(v) for v in sources]
              options.append("Add new")
              ADD_NEW = len(options) - 1
              action = userselect(
                  "Edit or Add %s:" % caption,
                  options,
                  default = ADD_NEW)
              if action == ADD_NEW:
                  sources.append(constructor(action))
              sources[action].interactive()
              add_more = userinput("Another?", default=True, inp=yn_inp, out=yn_out)
              
      def _i_config_target(self):
          self.target.interactive()
          return 6
          
      def _i_config_mode(self):
          self.mode.interactive()
          return 7

      def show_settings(self):
          """
          Display settings for the user.
          """
          print "Video Sources:"
          for video_source in self.video_sources:
              print "\t" + str(video_source)
          print "Audio Sources:"
          for audio_source in self.audio_sources:
              print "\t" + str(audio_source)
          print "Text (Subtitle) Sources:"
          for text_source in self.text_sources:
              print "\t" + str(text_source)
          print "Chapters:"
          for chapter in self.chapters:
              print "\t" + str(chapter)
          print "Output File:"
          print "\t" + str(self.target)
          print "Mode Settings:"
          print "\t" + str(self.mode_settings)
          return 8
                
      def start_conversion(self):
          """
          Spawn a process to run the conversion.
          """
          print "Not yet implemented"
          return 9
            
      def save_settings(self):
          """
          Save settings to disk.
          """
          print "Not yet implemented"
          return 10
            
      def load_settings(self):
          """
          Load settings from disk.
          """
          print "Not yet implemented"
          return 11
            
      def _exit(self):
          """
          End the program without doing anything.
          """
          print "Exiting without running conversion."
          sys.exit(0)
     

